package splitter;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.regex.Pattern;

/**
 * 基于 LangChain 的代码分割器实现
 */
public class LangChainCodeSplitter implements Splitter {
    private int chunkSize = 1000;
    private int chunkOverlap = 200;
    private static final Map<String, String> LANGUAGE_MAP = new HashMap<>();

    static {
        // 初始化语言映射
        LANGUAGE_MAP.put("javascript", "js");
        LANGUAGE_MAP.put("typescript", "js");
        LANGUAGE_MAP.put("python", "python");
        LANGUAGE_MAP.put("java", "java");
        LANGUAGE_MAP.put("cpp", "cpp");
        LANGUAGE_MAP.put("c++", "cpp");
        LANGUAGE_MAP.put("c", "cpp");
        LANGUAGE_MAP.put("go", "go");
        LANGUAGE_MAP.put("rust", "rust");
        LANGUAGE_MAP.put("php", "php");
        LANGUAGE_MAP.put("ruby", "ruby");
        LANGUAGE_MAP.put("swift", "swift");
        LANGUAGE_MAP.put("scala", "scala");
        LANGUAGE_MAP.put("html", "html");
        LANGUAGE_MAP.put("markdown", "markdown");
        LANGUAGE_MAP.put("md", "markdown");
        LANGUAGE_MAP.put("latex", "latex");
        LANGUAGE_MAP.put("tex", "latex");
        LANGUAGE_MAP.put("solidity", "sol");
        LANGUAGE_MAP.put("sol", "sol");
    }

    /**
     * 创建一个新的 LangChain 代码分割器
     */
    public LangChainCodeSplitter() {
    }

    /**
     * 创建一个新的 LangChain 代码分割器
     * 
     * @param chunkSize 代码块大小
     * @param chunkOverlap 代码块重叠大小
     */
    public LangChainCodeSplitter(int chunkSize, int chunkOverlap) {
        if (chunkSize > 0) this.chunkSize = chunkSize;
        if (chunkOverlap >= 0) this.chunkOverlap = chunkOverlap;
    }

    @Override
    public CompletableFuture<List<CodeChunk>> split(String code, String language, String filePath) throws Exception {
        return CompletableFuture.supplyAsync(() -> {
            try {
                // 获取映射后的语言
                String mappedLanguage = mapLanguage(language);
                if (mappedLanguage != null) {
                    // 使用特定语言的分割器
                    return splitByLanguage(code, language, filePath, mappedLanguage);
                } else {
                    // 使用通用分割器
                    return fallbackSplit(code, language, filePath);
                }
            } catch (Exception e) {
                System.err.println("Error splitting code: " + e.getMessage());
                e.printStackTrace();
                // 如果特定语言分割失败，使用通用分割器
                return fallbackSplit(code, language, filePath);
            }
        });
    }

    @Override
    public void setChunkSize(int chunkSize) {
        this.chunkSize = chunkSize;
    }

    @Override
    public void setChunkOverlap(int chunkOverlap) {
        this.chunkOverlap = chunkOverlap;
    }

    /**
     * 将语言名称映射到 LangChain 支持的格式
     * 
     * @param language 语言名称
     * @return 映射后的语言名称，如果不支持则返回 null
     */
    private String mapLanguage(String language) {
        return LANGUAGE_MAP.get(language.toLowerCase());
    }

    /**
     * 根据特定语言分割代码
     * 
     * @param code 代码内容
     * @param originalLanguage 原始语言名称
     * @param filePath 文件路径
     * @param mappedLanguage 映射后的语言名称
     * @return 代码块列表
     */
    private List<CodeChunk> splitByLanguage(String code, String originalLanguage, String filePath, String mappedLanguage) {
        // 这里应该使用特定语言的分割逻辑
        // 在 Java 中，我们需要实现自己的语言特定分割逻辑，而不是依赖 LangChain
        
        // 根据不同语言使用不同的分割策略
        switch (mappedLanguage) {
            case "js":
                return splitJavaScript(code, originalLanguage, filePath);
            case "python":
                return splitPython(code, originalLanguage, filePath);
            case "java":
                return splitJava(code, originalLanguage, filePath);
            default:
                // 对于其他语言，使用通用分割器
                return fallbackSplit(code, originalLanguage, filePath);
        }
    }

    /**
     * JavaScript 代码分割
     */
    private List<CodeChunk> splitJavaScript(String code, String language, String filePath) {
        // 简单的基于函数和类的分割
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        StringBuilder currentChunk = new StringBuilder();
        int startLine = 1;
        int currentLine = 1;
        
        Pattern functionPattern = Pattern.compile("^\\s*(function|class|const\\s+\\w+\\s*=\\s*\\(|\\w+\\s*=\\s*function)");
        
        for (String line : lines) {
            if (currentLine > 1 && functionPattern.matcher(line).find() && currentChunk.length() > 0) {
                // 找到新函数或类的开始，保存当前块
                chunks.add(new CodeChunk(
                    currentChunk.toString(),
                    startLine,
                    currentLine - 1,
                    language,
                    filePath
                ));
                
                // 开始新块
                currentChunk = new StringBuilder();
                startLine = currentLine;
            }
            
            currentChunk.append(line).append("\n");
            currentLine++;
            
            // 如果当前块太大，强制分割
            if (currentChunk.length() >= this.chunkSize) {
                chunks.add(new CodeChunk(
                    currentChunk.toString(),
                    startLine,
                    currentLine - 1,
                    language,
                    filePath
                ));
                
                currentChunk = new StringBuilder();
                startLine = currentLine;
            }
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                currentLine - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * Python 代码分割
     */
    private List<CodeChunk> splitPython(String code, String language, String filePath) {
        // 简单的基于函数和类的分割
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        StringBuilder currentChunk = new StringBuilder();
        int startLine = 1;
        int currentLine = 1;
        
        Pattern functionPattern = Pattern.compile("^\\s*(def|class)\\s+\\w+");
        
        for (String line : lines) {
            if (currentLine > 1 && functionPattern.matcher(line).find() && currentChunk.length() > 0) {
                // 找到新函数或类的开始，保存当前块
                chunks.add(new CodeChunk(
                    currentChunk.toString(),
                    startLine,
                    currentLine - 1,
                    language,
                    filePath
                ));
                
                // 开始新块
                currentChunk = new StringBuilder();
                startLine = currentLine;
            }
            
            currentChunk.append(line).append("\n");
            currentLine++;
            
            // 如果当前块太大，强制分割
            if (currentChunk.length() >= this.chunkSize) {
                chunks.add(new CodeChunk(
                    currentChunk.toString(),
                    startLine,
                    currentLine - 1,
                    language,
                    filePath
                ));
                
                currentChunk = new StringBuilder();
                startLine = currentLine;
            }
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                currentLine - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * Java 代码分割
     */
    private List<CodeChunk> splitJava(String code, String language, String filePath) {
        // 简单的基于方法和类的分割
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        StringBuilder currentChunk = new StringBuilder();
        int startLine = 1;
        int currentLine = 1;
        
        Pattern methodPattern = Pattern.compile("^\\s*(public|private|protected|static|\\s)*\\s*\\w+\\s+\\w+\\s*\\([^)]*\\)\\s*\\{?");
        Pattern classPattern = Pattern.compile("^\\s*(public|private|protected|static|\\s)*\\s*(class|interface|enum)\\s+\\w+");
        
        for (String line : lines) {
            if (currentLine > 1 && 
                (methodPattern.matcher(line).find() || classPattern.matcher(line).find()) && 
                currentChunk.length() > 0) {
                // 找到新方法或类的开始，保存当前块
                chunks.add(new CodeChunk(
                    currentChunk.toString(),
                    startLine,
                    currentLine - 1,
                    language,
                    filePath
                ));
                
                // 开始新块
                currentChunk = new StringBuilder();
                startLine = currentLine;
            }
            
            currentChunk.append(line).append("\n");
            currentLine++;
            
            // 如果当前块太大，强制分割
            if (currentChunk.length() >= this.chunkSize) {
                chunks.add(new CodeChunk(
                    currentChunk.toString(),
                    startLine,
                    currentLine - 1,
                    language,
                    filePath
                ));
                
                currentChunk = new StringBuilder();
                startLine = currentLine;
            }
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                currentLine - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }

    /**
     * 通用分割器作为后备
     * 
     * @param code 代码内容
     * @param language 语言
     * @param filePath 文件路径
     * @return 代码块列表
     */
    private List<CodeChunk> fallbackSplit(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        
        // 基于字符的递归分割
        List<String> textChunks = recursiveCharacterTextSplit(code, this.chunkSize, this.chunkOverlap);
        
        for (int i = 0; i < textChunks.size(); i++) {
            String chunk = textChunks.get(i);
            LineRange lines = estimateLines(chunk, code);
            
            chunks.add(new CodeChunk(
                chunk,
                lines.start,
                lines.end,
                language,
                filePath
            ));
        }
        
        return chunks;
    }

    /**
     * 基于字符的递归文本分割
     * 
     * @param text 文本内容
     * @param chunkSize 块大小
     * @param chunkOverlap 块重叠大小
     * @return 文本块列表
     */
    private List<String> recursiveCharacterTextSplit(String text, int chunkSize, int chunkOverlap) {
        List<String> chunks = new ArrayList<>();
        
        // 如果文本小于块大小，直接返回
        if (text.length() <= chunkSize) {
            chunks.add(text);
            return chunks;
        }
        
        // 尝试在分隔符处分割
        String[] separators = {"\n\n", "\n", ". ", ", ", " ", ""};
        
        for (String separator : separators) {
            if (separator.isEmpty()) {
                // 如果没有合适的分隔符，直接按字符分割
                for (int i = 0; i < text.length(); i += chunkSize - chunkOverlap) {
                    int end = Math.min(i + chunkSize, text.length());
                    chunks.add(text.substring(i, end));
                }
                return chunks;
            }
            
            List<String> splits = new ArrayList<>(Arrays.asList(text.split(Pattern.quote(separator))));
            
            // 过滤空字符串
            splits.removeIf(String::isEmpty);
            
            // 如果分割后的每个部分都小于块大小，使用此分隔符
            boolean goodSplit = true;
            for (String split : splits) {
                if (split.length() > chunkSize) {
                    goodSplit = false;
                    break;
                }
            }
            
            if (goodSplit) {
                // 合并小块，直到达到块大小
                StringBuilder currentChunk = new StringBuilder();
                
                for (String split : splits) {
                    if (currentChunk.length() + split.length() + separator.length() > chunkSize && currentChunk.length() > 0) {
                        chunks.add(currentChunk.toString());
                        currentChunk = new StringBuilder();
                    }
                    
                    if (currentChunk.length() > 0) {
                        currentChunk.append(separator);
                    }
                    
                    currentChunk.append(split);
                }
                
                if (currentChunk.length() > 0) {
                    chunks.add(currentChunk.toString());
                }
                
                return chunks;
            }
        }
        
        return chunks;
    }

    /**
     * 估计代码块的行范围
     */
    private static class LineRange {
        int start;
        int end;
        
        LineRange(int start, int end) {
            this.start = start;
            this.end = end;
        }
    }
    
    /**
     * 估计代码块在原始代码中的行范围
     * 
     * @param chunk 代码块
     * @param originalCode 原始代码
     * @return 行范围
     */
    private LineRange estimateLines(String chunk, String originalCode) {
        // 简单的行号估计
        String[] codeLines = originalCode.split("\n");
        String[] chunkLines = chunk.split("\n");
        
        // 在原始代码中查找块的位置
        int chunkStart = originalCode.indexOf(chunk);
        if (chunkStart == -1) {
            return new LineRange(1, chunkLines.length);
        }
        
        String beforeChunk = originalCode.substring(0, chunkStart);
        int startLine = beforeChunk.split("\n").length;
        int endLine = startLine + chunkLines.length - 1;
        
        return new LineRange(startLine, endLine);
    }
}